use std::pin::Pin;

use bytes::Bytes;
use tokio::select;
use tokio::sync::broadcast;
use tokio_stream::{Stream, StreamExt, StreamMap};
use tracing::debug;

use crate::{Connection, Db, Frame, Shutdown};
use crate::cmd::{Parse, ParseError};

#[derive(Debug)]
pub struct Subscribe {
    channels: Vec<String>,
}

type Messages = Pin<Box<dyn Stream<Item=Bytes> + Send>>;

impl Subscribe {
    pub(crate) fn new(channels: &[String]) -> Subscribe {
        Subscribe {
            channels: channels.to_vec(),
        }
    }

    pub(crate) fn parse_frames(parse: &mut Parse) -> crate::Result<Subscribe> {
        use ParseError::EndOfStream;
        // channels
        let mut channels = vec![parse.next_string()?];
        loop {
            match parse.next_string() {
                // 多个channel
                Ok(s) => channels.push(s),
                Err(EndOfStream) => break,
                Err(err) => return Err(err.into()),
            }
        }

        Ok(Subscribe { channels })
    }

    pub(crate) async fn apply(mut self, db: &Db, dst: &mut Connection, shutdown: &mut Shutdown) -> crate::Result<()> {
        let mut subscriptions = StreamMap::new();

        loop {
            for channel_name in self.channels.drain(..) {
                // 订阅
                subscribe_to_channel(channel_name, &mut subscriptions, db, dst).await?;
            }

            select! {
                // 写出消息
                Some((channel_name, msg)) = subscriptions.next() =>{
                    dst.write_frame(&make_message_frame(channel_name, msg)).await?;
                }
                _=shutdown.recv() => {
                    return Ok(());
                }
            }
        }
    }

    pub(crate) fn into_frame(self) -> Frame {
        let mut frame = Frame::array();
        frame.push_bulk(Bytes::from_static("subscribe".as_bytes()));
        for channel in self.channels {
            frame.push_bulk(Bytes::from(channel.into_bytes()));
        }
        frame
    }
}

/// 订阅
async fn subscribe_to_channel(
    channel_name: String,
    subscriptions: &mut StreamMap<String, Messages>,
    db: &Db,
    dst: &mut Connection,
) -> crate::Result<()> {
    // 从db 订阅获取broadcast::Receiver
    let mut rx = db.subscribe(channel_name.clone());
    // async stream
    let rx = Box::pin(async_stream::stream! {
        loop{
            match rx.recv().await{
                Ok(msg) => yield msg,
                Err(broadcast::error::RecvError::Lagged(_)) =>{}
                Err(_) => break,
            }
        }
    });
    // 存储stream
    subscriptions.insert(channel_name.clone(), rx);
    // 相应
    let response = make_subscribe_frame(channel_name, subscriptions.len());
    debug!(?response);
    dst.write_frame(&response).await?;

    Ok(())
}

fn make_subscribe_frame(channel_name: String, num_subs: usize) -> Frame {
    let mut response = Frame::array();
    response.push_bulk(Bytes::from_static(b"subscribe"));
    response.push_bulk(Bytes::from(channel_name));
    response.push_int(num_subs as u64);
    response
}

fn make_message_frame(channel_name: String, msg: Bytes) -> Frame {
    let mut response = Frame::array();
    response.push_bulk(Bytes::from_static(b"message"));
    response.push_bulk(Bytes::from(channel_name));
    response.push_bulk(msg);
    response
}